"""
获取并保存上证和深证股票代码的脚本
"""

import time
from struct import unpack

import h5py
import akshare as ak
import numpy as np
import pandas as pd
import os
import os.path as osp
os.environ['KMP_DUPLICATE_LIB_OK']='True'
from datetime import datetime

DataDir='E:\\tdx_data\data\\vipdoc'

CsvDataDir='E:\\tdx_data\\csvdata'
# E:\tdx_data\data\vipdoc\sh\lday
def day2csv(source_dir, file_name, target_dir):
    # 以二进制方式打开源文件
    ## 如果file_name 包含 sh
    if file_name.startswith('sh'):
         data_dir = osp.join(source_dir, 'sh', 'lday')
    else:
        data_dir = osp.join(source_dir, 'data', 'stock_zh_index_daily_hfq')
    source_file = open(data_dir + os.sep + file_name, 'rb')
    buf = source_file.read()
    source_file.close()

    # 打开目标文件
    target_file = open(target_dir + os.sep + file_name[: file_name.rindex('.')] + '.csv', 'w')
    buf_size = len(buf)
    rec_count = int(buf_size / 32)
    begin = 0
    end = 32
    header = str('date') + ',' + str('open') + ',' + str('high') + ',' + str('low') + ',' \
             + str('close') + ',' + str('amount') + ',' + str('volume') + '\n'
    target_file.write(header)
    for i in range(rec_count):
        a = unpack('IIIIIfII', buf[begin:end])
        # 处理date数据
        year = a[0] // 10000
        month = (a[0] % 10000) // 100
        day = (a[0] % 10000) % 100
        date = '{}-{:02d}-{:02d}'.format(year, month, day)

        line = date + ',' + str(a[1] / 100.0) + ',' + str(a[2] / 100.0) + ',' \
               + str(a[3] / 100.0) + ',' + str(a[4] / 100.0) + ',' + str(a[5]) + ',' \
               + str(a[6]) + '\n'
        target_file.write(line)
        begin += 32
        end += 32
    target_file.close()
def main():

    day2csv(DataDir, 'sh000001.day', CsvDataDir)



if __name__ == "__main__":
    main()